#coding=utf-8
#! /bin/env python

"""This program convert cif file format into msd file format
mainly by MadLee Hawkweed did some little improvement
@2009-01-06
note need pykuai"""

import sys
import string
from dffpy import *
from math import radians, degrees
from typeopls import TypingOPLS

USAGE = """usage: cif2msd.py file.cif file.msd"""

atom2index = {}

def parse_float(token):
    result = ""
    for i in token:
        if i == "+" or i == "-" or i == "." or i == "e" or i == "E" or i in string.digits:
            result += i
    return float(result)
    
def parse_pbc(file, line):
    tokens = line.split()
    a = parse_float(tokens[1])
    b = parse_float(file.next().split()[1])
    c = parse_float(file.next().split()[1])
    alpha = parse_float(file.next().split()[1])
    beta = parse_float(file.next().split()[1])
    gamma = parse_float(file.next().split()[1])
    return PBC(a, b, c, radians(alpha), radians(beta), radians(gamma))
    
def parse_atoms(file):
    result = []
    for line in file:
        if line[0] != '_':
            tokens = line.split()
            if len(tokens) >= 5:
                a = Atom(tokens[1])
                a.coord.x = float(tokens[2])
                a.coord.y = float(tokens[3])
                a.coord.z = float(tokens[4])
                result.append((tokens[0], a))
            else:
                break
    return result

def parse_bonds(file, atom2index):
    result = []
    for line in file:
        if line[0] != '_':
            tokens = line.split()
            if len(tokens) >= 2:
                a1 = atom2index[tokens[0]]
                a2 = atom2index[tokens[1]]
                result.append(Bond(a1, a2))
            else:
                break
    return result

def savemsd(filename, pbc, atoms, bonds):
    mol = MolecularModel.setup(atoms, bonds)
    mol.pbc = pbc
    TypingOPLS.instance.assign_types(mol)
    mol.save(filename)
    
    # with open(filename, 'wb') as file:
    #     file.write("#DFF:MSD\n")
    #     file.write("#Model Structure Data File    Energy = 0.0\n")
    #     file.write("PBC: %12.4f %12.4f %12.4f %12.4f %12.4f %12.4f\n"%(pbc.a, pbc.b, pbc.c, degrees(pbc.alpha), degrees(pbc.beta), degrees(pbc.gamma)))
    #     file.write('%6d\n'%len(atoms))
    #     for i in range(len(atoms)):
    #         atom = atoms[i][1]
    #         coord = pbc.va * atom.coord.x + pbc.vb * atom.coord.y + pbc.vc * atom.coord.z
    #         atom.coord = coord
    #         line = "%6d %8s %3d ? ? 0.0 %12.4f %12.4f %12.4f  1 RES   0\n" % (i+1, atoms[i][0], atom.number, coord.x, coord.y, coord.z)
    #         file.write(line)
    #     file.write('%6d\n'%len(bonds))
    #     for i in bonds:
    #         line = "%6s %6s     1\n"%(i[0],i[1])
    #         file.write(line)

if __name__ == '__main__':
    if len(sys.argv) != 3:
        print USAGE
    else:  
        with open(sys.argv[1]) as file:
            for line in file:
                if line.startswith("_atom_site_label"):
                    atoms = parse_atoms(file)
                    atom2index = {}
                    ii = 1
                    for i in atoms:
                        atom2index[i[0]] = ii
                        ii += 1
                elif line.startswith("_geom_bond_atom_site_label_1"):
                    bonds = parse_bonds(file, atom2index)
                elif line.startswith("_cell_length_a"):
                    pbc = parse_pbc(file, line)
                    
        savemsd(sys.argv[2], pbc, atoms, bonds)    
            
